import torch
from dataclasses import MISSING
from isaaclab.managers.action_manager import ActionTerm, ActionTermCfg
from isaaclab.envs.mdp.actions.joint_actions import JointAction
from isaaclab.utils import configclass  # 导入 configclass 装饰器
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedEnv

class PDController:
    def __init__(self, kp: float, kd: float):
        """
        初始化新的 PD 控制器。

        :param kp: 比例增益
        :param kd: 微分增益
        """
        self.kp = kp
        self.kd = kd

    def compute_torque(self, current_pos: torch.Tensor, target_pos: torch.Tensor, current_vel: torch.Tensor) -> torch.Tensor:
        """
        根据当前位置、目标位置和当前速度计算力矩。

        :param current_pos: 当前位置
        :param target_pos: 目标位置
        :param current_vel: 当前速度
        :return: 计算得到的力矩
        """
        pos_error = target_pos - current_pos
        torque = self.kp * pos_error - self.kd * current_vel
        return torque

class PDJointAction(JointAction):
    """新的关节动作项，使用 PD 控制器将位置动作转换为力矩命令。"""

    cfg: 'PDJointActionCfg'  # 使用字符串形式的类型注解
    """动作项的配置。"""
    _pd_controller: PDController
    """用于将位置转换为力矩的 PD 控制器。"""

    def __init__(self, cfg: 'PDJointActionCfg', env: 'ManagerBasedEnv'):  # 使用字符串形式的类型注解
        super().__init__(cfg, env)
        # 初始化新的 PD 控制器
        self._pd_controller = PDController(cfg.kp, cfg.kd)

    def apply_actions(self):
        # 获取当前关节位置和速度
        current_pos = self._asset.data.joint_pos[:, self._joint_ids]
        current_vel = self._asset.data.joint_vel[:, self._joint_ids]
        # 将 processed_actions 作为目标位置
        target_pos = self.processed_actions
        # 使用新的 PD 控制器计算力矩
        torque = self._pd_controller.compute_torque(current_pos, target_pos, current_vel)
        # 设置关节力矩目标
        self._asset.set_joint_effort_target(torque, joint_ids=self._joint_ids)

@configclass
class PDJointActionCfg(ActionTermCfg):
    """新的关节位置到力矩动作项配置，使用 PD 控制器。"""

    class_type: type[ActionTerm] = PDJointAction

    joint_names: list[str] = MISSING
    """动作将映射到的关节名称列表或正则表达式列表。"""

    scale: float | dict[str, float] = 1.0
    """动作的缩放因子（浮点数或正则表达式字典）。默认为 1.0。"""

    offset: float | dict[str, float] = 0.0
    """动作的偏移因子（浮点数或正则表达式字典）。默认为 0.0。"""

    clip: dict[str, list[float]] | None = None
    """动作的裁剪因子（正则表达式字典）。默认为 None。"""

    kp: float = MISSING
    """PD 控制器的比例增益。"""

    kd: float = MISSING
    """PD 控制器的微分增益。"""

    preserve_order: bool = False
    """是否保留解析后关节的顺序。默认为 False。"""